# Copyright (c) 2024-2024 Huawei Technologies Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from components.utils.parser import BaseCommand
from components.utils.security_check import check_input_path_legality, check_output_path_legality
from components.utils.log import set_log_level, LOG_LEVELS_LOWER
from msit_opcheck.opchecker import OpChecker
from msit_opcheck.autofuse.fusion_op_check import FuseOpChecker


class OpcheckCommand(BaseCommand):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.parser = None

    def add_arguments(self, parser, **kwargs):
        parser.add_argument(
            '--input',
            '-i',
            required=True,
            dest='input',
            type=check_input_path_legality,
            help='Configure the GE dump data path. In "single" mode, you can set it to the path '
                 'generated by the "msit debug dump" tool.')
        
        parser.add_argument(
            '--mode',
            '-m',
            required=False,
            dest='mode',
            default='single',
            choices=['single', 'autofuse'],
            help='Configure the opcheck scenario. Currently, two modes are supported: "single" and "autofuse".')
        # 当前msit debug dump对于自动融合场景不适配，因此需要手动指定graph路径
        parser.add_argument(
            '--graph-path',
            '-gp',
            required=False,
            dest='graph_path',
            type=check_input_path_legality,
            help='Configure the GE dump graph path.')

        parser.add_argument(
            '--output',
            '-o',
            required=False,
            dest='output',
            type=check_output_path_legality,
            default='./',
            help='Result data save directory. For example:--output /tmp/output.')
        
        parser.add_argument("--log-level", "-l", default="info", choices=LOG_LEVELS_LOWER, help="specify log level.")

    def handle(self, args):
        set_log_level(args.log_level)
        if args.mode == "single":
            opchecker = OpChecker(args)
        else:
            if not args.graph_path:
                raise ValueError("If the parameter '--mode' is set to 'autofuse', \
                                  it must be used together with parameter '--graph-path'.")
            opchecker = FuseOpChecker(args)
        opchecker.start_test()


def get_opcheck_cmd_ins():
    help_info = "Operation check tool for GE compile model."
    opcheck_instance = OpcheckCommand("opcheck", help_info)
    return opcheck_instance